package org.bdware.irp.crypto;

import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.*;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.KeyType;
import org.apache.log4j.Logger;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.text.ParseException;

public class CertUtils {

    static Logger logger = Logger.getLogger(org.bdware.irp.crypto.CertUtils.class);

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static String Sign(byte[] data, JWK jwk) throws Exception {
        JWSSigner jwsSigner;
        JWSObject jwsObject;

        if(jwk.getKeyType() == KeyType.RSA){
            jwsSigner = new RSASSASigner(jwk.toRSAKey());
            jwsObject = new JWSObject(
                    new JWSHeader.Builder(JWSAlgorithm.RS256).keyID(jwk.getKeyID()).build(),
                    new Payload(data));
        }else if(jwk.getKeyType() == KeyType.EC){
            jwsSigner = new ECDSASigner(jwk.toECKey());
            // Creates the JWS object with payload
            jwsObject = new JWSObject(
                    new JWSHeader.Builder(JWSAlgorithm.ES256).keyID(jwk.getKeyID()).build(),
                    new Payload(data));
        }else{
            logger.error("unsupported Algorithm");
            throw new NoSuchAlgorithmException("unsupported Algorithm");
        }
        // Compute the EC signature
        jwsObject.sign(jwsSigner);
        // Serialize the JWS to compact form
        String s = jwsObject.serialize(true);
        return s;
    }

    public static boolean verify(byte[] payload, String jws, JWK pkToVerify) throws Exception {
        JWSObject signature = JWSObject.parse(jws,new Payload(payload));
        JWSVerifier verifier;
        if(pkToVerify.getKeyType() == KeyType.RSA){
            verifier = new RSASSAVerifier(pkToVerify.toRSAKey().toRSAPublicKey());
        }else if(pkToVerify.getKeyType() == KeyType.EC){
            verifier = new ECDSAVerifier(pkToVerify.toECKey().toECPublicKey());
        }else{
            logger.error("unsupported Algorithm");
            return false;
        }
        return signature.verify(verifier);
    }

    public static byte[] encrypt(byte[] payload, JWK pkToEncrypt) throws JOSEException {
        logger.debug("===do encryption");
        JWEAlgorithm alg;
        EncryptionMethod enc = EncryptionMethod.A256GCM;
        JWEObject jweObject;
        if(pkToEncrypt.getKeyType() == KeyType.RSA){
            alg = JWEAlgorithm.RSA_OAEP_256;
            jweObject= new JWEObject(
                    new JWEHeader(alg,enc),
                    new Payload(payload)
            );
            jweObject.encrypt(new RSAEncrypter(pkToEncrypt.toRSAKey().toRSAPublicKey()));
        }else if(pkToEncrypt.getKeyType() == KeyType.EC){
            alg = JWEAlgorithm.ECDH_ES;
            jweObject= new JWEObject(
                    new JWEHeader(alg,enc),
                    new Payload(payload)
            );
            jweObject.encrypt(new ECDHEncrypter(pkToEncrypt.toECKey().toECPublicKey()));
        }else{
            throw new JOSEException("only support RSA or EC key");
        }
        return jweObject.serialize().getBytes();
    }

    public static byte[] decrypt(byte[] cipherText, JWK skToDecrypt) throws ParseException, JOSEException {
        logger.debug("===do decryption");
        JWEObject decrypt = JWEObject.parse(new String(cipherText));
        if(skToDecrypt.getKeyType() == KeyType.RSA){
            decrypt.decrypt(new RSADecrypter(skToDecrypt.toRSAKey().toRSAPrivateKey()));
        }else if(skToDecrypt.getKeyType() == KeyType.EC){
            decrypt.decrypt(new ECDHDecrypter(skToDecrypt.toECKey().toECPrivateKey()));
        }else{
            throw new JOSEException("only support RSA or EC key");
        }
        return decrypt.getPayload().toBytes();
    }
}
